// Copyright Epic Games, Inc. All Rights Reserved.

#include "admin_cmd.h"
#include <zencore/filesystem.h>
#include <zencore/logging.h>
#include <zenhttp/formatters.h>
#include <zenhttp/httpclient.h>
#include <zenhttp/httpcommon.h>
#include <zenutil/basicfile.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <cpr/cpr.h>
ZEN_THIRD_PARTY_INCLUDES_END

using namespace std::literals;

namespace zen {

ScrubCommand::ScrubCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "n", "dry", "Dry run (do not delete any data)", cxxopts::value(m_DryRun), "<bool>");
	m_Options.add_option("", "", "no-gc", "Do not perform GC after scrub pass", cxxopts::value(m_NoGc), "<bool>");
	m_Options.add_option("", "", "no-cas", "Do not scrub CAS stores", cxxopts::value(m_NoCas), "<bool>");
}

ScrubCommand::~ScrubCommand() = default;

int
ScrubCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	HttpClient Http(m_HostName);

	HttpClient::KeyValueMap Params{{"skipdelete", ToString(m_DryRun)}, {"skipgc", ToString(m_NoGc)}, {"skipcid", ToString(m_NoCas)}};

	if (HttpClient::Response Response = Http.Post("/admin/scrub"sv, /* headers */ HttpClient::KeyValueMap{}, Params))
	{
		ZEN_CONSOLE("scrub started OK: {}", Response.ToText());

		return 0;
	}
	else if (int StatusCode = (int)Response.StatusCode)
	{
		ZEN_ERROR("scrub start failed: {}: {} ({})",
				  (int)Response.StatusCode,
				  ReasonStringForHttpResultCode((int)Response.StatusCode),
				  Response.ToText());
	}
	else
	{
		ZEN_ERROR("scrub start failed: {}", Response.ToText());
	}

	return 1;
}

//////////////////////////////////////////////////////////////////////////

GcCommand::GcCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("",
						 "s",
						 "smallobjects",
						 "Collect small objects",
						 cxxopts::value(m_SmallObjects)->default_value("false"),
						 "<smallobjects>");
	m_Options.add_option("", "", "skipcid", "Skip collection of CAS data", cxxopts::value(m_SkipCid)->default_value("false"), "<skipcid>");
	m_Options.add_option("",
						 "n",
						 "skipdelete",
						 "Skip deletion of data (dryrun)",
						 cxxopts::value(m_SkipDelete)->default_value("false"),
						 "<skipdelete>");
	m_Options.add_option("",
						 "m",
						 "maxcacheduration",
						 "Max cache lifetime (in seconds)",
						 cxxopts::value(m_MaxCacheDuration)->default_value("0"),
						 "<maxcacheduration>");
	m_Options.add_option("",
						 "d",
						 "disksizesoftlimit",
						 "Max disk usage size (in bytes)",
						 cxxopts::value(m_DiskSizeSoftLimit)->default_value("0"),
						 "<disksizesoftlimit>");
	m_Options
		.add_option("", "", "usegcv1", "Force use of GC version 1", cxxopts::value(m_ForceUseGCV1)->default_value("false"), "<usegcv2>");
	m_Options
		.add_option("", "", "usegcv2", "Force use of GC version 2", cxxopts::value(m_ForceUseGCV2)->default_value("false"), "<usegcv2>");
	m_Options.add_option("",
						 "",
						 "compactblockthreshold",
						 "How much of a compact block should be used to skip compacting the block. 0 - compact only empty eligible blocks, "
						 "100 - compact all non-full eligible blocks.",
						 cxxopts::value(m_CompactBlockThreshold)->default_value("60"),
						 "<compactblockthreshold>");
	m_Options
		.add_option("", "", "verbose", "Enable verbose logging for GC", cxxopts::value(m_Verbose)->default_value("false"), "<verbose>");
	m_Options.add_option("",
						 "",
						 "single-threaded",
						 "Force GC to run single threaded",
						 cxxopts::value(m_SingleThreaded)->default_value("false"),
						 "<single-threaded>");
}

GcCommand::~GcCommand()
{
}

int
GcCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	cpr::Parameters Params;
	Params.Add({"smallobjects", m_SmallObjects ? "true" : "false"});
	if (m_MaxCacheDuration != 0)
	{
		Params.Add({"maxcacheduration", fmt::format("{}", m_MaxCacheDuration)});
	}
	if (m_DiskSizeSoftLimit != 0)
	{
		Params.Add({"disksizesoftlimit", fmt::format("{}", m_DiskSizeSoftLimit)});
	}
	Params.Add({"skipcid", m_SkipCid ? "true" : "false"});
	Params.Add({"skipdelete", m_SkipDelete ? "true" : "false"});
	if (m_ForceUseGCV1)
	{
		if (m_ForceUseGCV2)
		{
			throw OptionParseException("only usegcv1 or usegcv2 can be selected, not both");
		}
		Params.Add({"forceusegcv1", "true"});
	}
	if (m_ForceUseGCV2)
	{
		Params.Add({"forceusegcv2", "true"});
	}
	if (m_CompactBlockThreshold)
	{
		Params.Add({"compactblockthreshold", fmt::format("{}", m_CompactBlockThreshold)});
	}
	Params.Add({"verbose", m_Verbose ? "true" : "false"});
	Params.Add({"singlethreaded", m_SingleThreaded ? "true" : "false"});

	cpr::Session Session;
	Session.SetHeader(cpr::Header{{"Accept", "application/json"}});
	Session.SetUrl({fmt::format("{}/admin/gc", m_HostName)});
	Session.SetParameters(Params);

	cpr::Response Result = Session.Post();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("OK: {}", Result.text);
		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("GC start failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("GC start failed: {}", Result.error.message);
	}

	return 1;
}

GcStatusCommand::GcStatusCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "d", "details", "Show detailed GC report", cxxopts::value(m_Details)->default_value("false"), "<details>");
}

GcStatusCommand::~GcStatusCommand()
{
}

int
GcStatusCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	cpr::Session Session;
	Session.SetHeader(cpr::Header{{"Accept", "application/json"}});
	Session.SetUrl({fmt::format("{}/admin/gc", m_HostName)});
	if (m_Details)
	{
		Session.SetParameters({{"details", "true"}});
	}

	cpr::Response Result = Session.Get();

	if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("OK: {}", Result.text);
		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("GC status failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("GC status failed: {}", Result.error.message);
	}

	return 1;
}

GcStopCommand::GcStopCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
}

GcStopCommand::~GcStopCommand()
{
}

int
GcStopCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	cpr::Session Session;
	Session.SetUrl({fmt::format("{}/admin/gc-stop", m_HostName)});
	cpr::Response Result = Session.Post();

	if (static_cast<HttpResponseCode>(Result.status_code) == HttpResponseCode::Accepted)
	{
		ZEN_CONSOLE("OK: {}", "Cancel request accepted");
		return 0;
	}
	else if (zen::IsHttpSuccessCode(Result.status_code))
	{
		ZEN_CONSOLE("OK: {}", "No GC running");
		return 0;
	}

	if (Result.status_code)
	{
		ZEN_ERROR("GC status failed: {}: {} ({})", Result.status_code, Result.reason, Result.text);
	}
	else
	{
		ZEN_ERROR("GC status failed: {}", Result.error.message);
	}

	return 1;
}

////////////////////////////////////////////

JobCommand::JobCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "j", "jobid", "Job id", cxxopts::value(m_JobId), "<jobid>");
	m_Options.add_option("", "c", "cancel", "Cancel job id", cxxopts::value(m_Cancel), "<cancel>");
}

JobCommand::~JobCommand() = default;

int
JobCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	using namespace std::literals;

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	HttpClient Http(m_HostName);

	if (m_Cancel)
	{
		if (m_JobId == 0)
		{
			ZEN_ERROR("Job id must be given");
			return 1;
		}
	}
	std::string Url = m_JobId != 0 ? fmt::format("/admin/jobs/{}", m_JobId) : "/admin/jobs";

	if (m_Cancel)
	{
		if (HttpClient::Response Result = Http.Delete(Url, HttpClient::Accept(ZenContentType::kJSON)))
		{
			ZEN_CONSOLE("{}", Result);
		}
		else
		{
			Result.ThrowError("failed cancelling job"sv);
			return 1;
		}
	}
	else if (HttpClient::Response Result = Http.Get(Url, HttpClient::Accept(ZenContentType::kJSON)))
	{
		ZEN_CONSOLE("{}", Result.ToText());
	}
	else
	{
		Result.ThrowError("failed fetching job info"sv);
		return 1;
	}

	return 0;
}

////////////////////////////////////////////

LoggingCommand::LoggingCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
	m_Options.add_option("", "", "cache-write-log", "Enable cache write logging", cxxopts::value(m_CacheWriteLog), "<enable/disable>");
	m_Options.add_option("", "", "cache-access-log", "Enable cache access logging", cxxopts::value(m_CacheAccessLog), "<enable/disable>");
	m_Options
		.add_option("", "", "set-log-level", "Set zenserver log level", cxxopts::value(m_SetLogLevel), "<trace/debug/info/warning/error>");
	m_Options.add_option("",
						 "",
						 "copy-log",
						 "Copy the server log file from a local zenserver instance",
						 cxxopts::value(m_ServerLogTarget),
						 "<targetpath>");
	m_Options.add_option("",
						 "",
						 "copy-cache-log",
						 "Copy the server cache log file from a local zenserver instance",
						 cxxopts::value(m_CacheLogTarget),
						 "<targetpath>");
	m_Options.add_option("",
						 "",
						 "copy-http-log",
						 "Copy the server http log file from a local zenserver instance",
						 cxxopts::value(m_HttpLogTarget),
						 "<targetpath>");
}

LoggingCommand::~LoggingCommand() = default;

int
LoggingCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	using namespace std::literals;

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	HttpClient Http(m_HostName);

	HttpClient::KeyValueMap Parameters;

	if (!m_CacheWriteLog.empty())
	{
		if (m_CacheWriteLog == "enable")
		{
			(*Parameters)["cacheenablewritelog"] = "true";
		}
		else if (m_CacheWriteLog == "disable")
		{
			(*Parameters)["cacheenablewritelog"] = "false";
		}
		else
		{
			ZEN_ERROR("Invalid value for parameter 'cache-write-log'. Use 'enable' or 'disable'");
			return 1;
		}
	}

	if (!m_CacheAccessLog.empty())
	{
		if (m_CacheAccessLog == "enable")
		{
			(*Parameters)["cacheenableaccesslog"] = "true";
		}
		else if (m_CacheAccessLog == "disable")
		{
			(*Parameters)["cacheenableaccesslog"] = "false";
		}
		else
		{
			ZEN_ERROR("Invalid value for parameter 'cache-access-log'. Use 'enable' or 'disable'");
			return 1;
		}
	}

	if (!m_SetLogLevel.empty())
	{
		(*Parameters)["loglevel"] = m_SetLogLevel;
	}

	if ((*Parameters).empty())
	{
		if (HttpClient::Response Result = Http.Get("/admin/logs", HttpClient::Accept(ZenContentType::kCbObject)))
		{
			ZEN_CONSOLE("{}", Result.ToText());
			const CbObject LogsResponse = Result.AsObject();

			auto CopyLog = [](std::string_view SourceName, std::string_view SourcePath, std::string_view TargetPath) -> bool {
				if (SourcePath.empty())
				{
					ZEN_ERROR("Failed to retrieve {} log path", SourceName);
					return false;
				}
				if (!CopyFile(SourcePath, TargetPath, {}))
				{
					ZEN_ERROR("Failed to copy {} log file {} to output file '{}'", SourceName, SourcePath, TargetPath);
					return false;
				}
				return true;
			};

			if (!m_ServerLogTarget.empty())
			{
				if (!CopyLog("server", LogsResponse["Logfile"].AsString(), m_ServerLogTarget))
				{
					return 1;
				}
			}

			if (!m_CacheLogTarget.empty())
			{
				if (!CopyLog("cache", LogsResponse["cache"].AsObjectView()["Logfile"].AsString(), m_CacheLogTarget))
				{
					return 1;
				}
			}
			if (!m_HttpLogTarget.empty())
			{
				if (!CopyLog("http", LogsResponse["http"].AsObjectView()["Logfile"].AsString(), m_HttpLogTarget))
				{
					return 1;
				}
			}
		}
		else
		{
			Result.ThrowError("failed fetching log info"sv);
			return 1;
		}
		return 0;
	}
	if (HttpClient::Response Result = Http.Post("/admin/logs", HttpClient::KeyValueMap{}, Parameters))
	{
		ZEN_CONSOLE("{}", Result.ToText());
	}
	else
	{
		Result.ThrowError("failed setting log info"sv);
		return 1;
	}

	return 0;
}

//////////////////////////////////////////////////////////////////////////

FlushCommand::FlushCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value(""), "<hosturl>");
}

FlushCommand::~FlushCommand() = default;

int
FlushCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	m_HostName = ResolveTargetHostSpec(m_HostName);

	if (m_HostName.empty())
	{
		throw OptionParseException("unable to resolve server specification");
	}

	zen::HttpClient Http(m_HostName);

	if (zen::HttpClient::Response Response = Http.Post("/admin/flush"sv))
	{
		ZEN_CONSOLE("OK: {}", Response.ToText());

		return 0;
	}
	else if (int StatusCode = (int)Response.StatusCode)
	{
		ZEN_ERROR("flush failed: {}: {} ({})",
				  (int)Response.StatusCode,
				  ReasonStringForHttpResultCode((int)Response.StatusCode),
				  Response.ToText());
	}
	else
	{
		ZEN_ERROR("flush failed: {}", Response.ToText());
	}

	return 1;
}

//////////////////////////////////////////////////////////////////////////

CopyStateCommand::CopyStateCommand()
{
	m_Options.add_options()("h,help", "Print help");
	m_Options.add_option("", "", "data-path", "Zen server source data path", cxxopts::value(m_DataPath), "<data-path>");
	m_Options.add_option("", "", "target-path", "Target data path", cxxopts::value(m_TargetPath), "<target-path>");
	m_Options.add_option("",
						 "",
						 "skip-logs",
						 "Only copy state index files, not log files (recommended to issue a zen flush before using this option)",
						 cxxopts::value(m_SkipLogs)->default_value("false"),
						 "<skip-logs>");
	m_Options.parse_positional({"data-path", "target-path"});
}

CopyStateCommand::~CopyStateCommand() = default;

static void
Copy(const std::filesystem::path& Source, const std::filesystem::path& Target)
{
	CreateDirectories(Target.parent_path());

	CopyFileOptions Options;
	CopyFile(Source, Target, Options);
}

static bool
TryCopy(const std::filesystem::path& Source, const std::filesystem::path& Target)
{
	if (!std::filesystem::is_regular_file(Source))
	{
		return false;
	}

	CreateDirectories(Target.parent_path());

	CopyFileOptions Options;
	return CopyFile(Source, Target, Options);
}

int
CopyStateCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	if (m_DataPath.empty())
	{
		throw OptionParseException("data path must be given");
	}

	if (!std::filesystem::is_directory(m_DataPath))
	{
		throw OptionParseException("data path must exist");
	}

	if (m_TargetPath.empty())
	{
		throw OptionParseException("target path must be given");
	}

	std::filesystem::path RootManifestPath		 = m_DataPath / "root_manifest";
	std::filesystem::path TargetRootManifestPath = m_TargetPath / "root_manifest";

	if (!TryCopy(RootManifestPath, TargetRootManifestPath))
	{
		throw OptionParseException("data path is invalid, missing root_manifest");
	}

	std::filesystem::path CachePath		  = m_DataPath / "cache";
	std::filesystem::path TargetCachePath = m_TargetPath / "cache";

	// Copy cache state
	DirectoryContent CacheDirectoryContent;
	GetDirectoryContent(CachePath, DirectoryContent::IncludeDirsFlag, CacheDirectoryContent);
	for (const std::filesystem::path& NamespacePath : CacheDirectoryContent.Directories)
	{
		std::filesystem::path NamespaceName		  = NamespacePath.filename();
		std::filesystem::path TargetNamespacePath = TargetCachePath / NamespaceName;

		DirectoryContent CacheNamespaceContent;
		GetDirectoryContent(NamespacePath, DirectoryContent::IncludeDirsFlag, CacheNamespaceContent);

		for (const std::filesystem::path& BucketPath : CacheNamespaceContent.Directories)
		{
			std::filesystem::path BucketName	   = BucketPath.filename();
			std::filesystem::path TargetBucketPath = TargetNamespacePath / BucketName;

			// TODO: make these use file naming helpers from cache implementation?

			std::filesystem::path ManifestPath		 = BucketPath / "zen_manifest";
			std::filesystem::path TargetManifestPath = TargetBucketPath / "zen_manifest";
			if (TryCopy(ManifestPath, TargetManifestPath))
			{
				if (!m_SkipLogs)
				{
					std::filesystem::path LogName		= fmt::format("{}.{}", BucketName.string(), "slog");
					std::filesystem::path LogPath		= BucketPath / LogName;
					std::filesystem::path TargetLogPath = TargetBucketPath / LogName;
					Copy(LogPath, TargetLogPath);
				}

				std::filesystem::path IndexName		  = fmt::format("{}.{}", BucketName.string(), "uidx");
				std::filesystem::path IndexPath		  = BucketPath / IndexName;
				std::filesystem::path TargetIndexPath = TargetBucketPath / IndexName;
				TryCopy(IndexPath, TargetIndexPath);

				std::filesystem::path MetaName		 = fmt::format("{}.{}", BucketName.string(), "meta");
				std::filesystem::path MetaPath		 = BucketPath / MetaName;
				std::filesystem::path TargetMetaPath = TargetBucketPath / MetaName;
				TryCopy(MetaPath, TargetMetaPath);
			}
		}
	}

	std::filesystem::path CasPath		= m_DataPath / "cas";
	std::filesystem::path TargetCasPath = m_TargetPath / "cas";

	{
		std::filesystem::path UCasRootPath		 = CasPath / ".ucas_root";
		std::filesystem::path TargetUCasRootPath = TargetCasPath / ".ucas_root";
		Copy(UCasRootPath, TargetUCasRootPath);
	}

	if (!m_SkipLogs)
	{
		std::filesystem::path LogPath		= CasPath / "cas.ulog";
		std::filesystem::path TargetLogPath = TargetCasPath / "cas.ulog";
		Copy(LogPath, TargetLogPath);
	}

	{
		std::filesystem::path IndexPath		  = CasPath / "cas.uidx";
		std::filesystem::path TargetIndexPath = TargetCasPath / "cas.uidx";
		TryCopy(IndexPath, TargetIndexPath);
	}

	{
		std::filesystem::path SobsRootPath		 = CasPath / "sobs";
		std::filesystem::path TargetSobsRootPath = TargetCasPath / "sobs";

		{
			std::filesystem::path SobsIndexPath		  = SobsRootPath / "sobs.uidx";
			std::filesystem::path TargetSobsIndexPath = TargetSobsRootPath / "sobs.uidx";
			TryCopy(SobsIndexPath, TargetSobsIndexPath);

			if (!m_SkipLogs)
			{
				std::filesystem::path SobsLogPath		= SobsRootPath / "sobs.ulog";
				std::filesystem::path TargetSobsLogPath = TargetSobsRootPath / "sobs.ulog";
				Copy(SobsLogPath, TargetSobsLogPath);
			}
		}
	}

	{
		std::filesystem::path TobsRootPath		 = CasPath / "tobs";
		std::filesystem::path TargetTobsRootPath = TargetCasPath / "tobs";
		{
			std::filesystem::path TobsIndexPath		  = TobsRootPath / "tobs.uidx";
			std::filesystem::path TargetTobsIndexPath = TargetTobsRootPath / "tobs.uidx";
			TryCopy(TobsIndexPath, TargetTobsIndexPath);

			if (!m_SkipLogs)
			{
				std::filesystem::path TobsLogPath		= TobsRootPath / "tobs.ulog";
				std::filesystem::path TargetTobsLogPath = TargetTobsRootPath / "tobs.ulog";
				Copy(TobsLogPath, TargetTobsLogPath);
			}
		}
	}

	return 0;
}

}  // namespace zen
